#include <cstring>
#include <stdexcept>
#include <vector>
#include <iostream>
#include <algorithm>
#include <cstdlib>
#include <cstdio>
#include <unordered_map>
#include <string>
#include <filesystem>
#include <opencv2/opencv.hpp>
#include <getopt.h>
#include "satlab.h"
#include <regex>
#include <fstream>
#include <chrono>

using Clock = std::chrono::high_resolution_clock;
using Duration = std::chrono::duration<double>;

void print_usage() {
    std::cout << "Usage:\n"
              << "  satlab gen-scope-image <source CNF file> <origin|sorted> <output image file>\n"
              << "  satlab gen-ascii-image <source CNF file> <min|max|mean|direct> <output image file>\n"
              << "  satlab gen-bin-scope-images <source CNF file> <origin|sorted> <output directory>\n"
              << "  satlab gen-block-images <source CNF file> <origin|sorted> <output directory>\n";
}

void run_gen_scope_image(const std::string& src_file, const std::string& type, const std::string& out_file) {
    if (type != "origin" && type != "sorted") {
        std::cerr << "[gen-scope-image] Type must be 'origin' or 'sorted'." << std::endl;
        print_usage();
        exit(1);
    }
    try {
        SATLAB lab(512, 512);
        lab.parse(src_file.c_str());
        if (type == "sorted") {
            lab.mat.sortVecs();
        }
        lab.updateAllStats();
        cv::Mat img;
        lab.createImage(img);
        cv::imwrite(out_file, img);
        std::cout << "[gen-scope-image] Processed: " << src_file << " -> " << out_file << std::endl;
    } catch (const std::exception& e) {
        std::cerr << "[gen-scope-image] Error: " << e.what() << std::endl;
        exit(2);
    }
}

void run_gen_ascii_image(const std::string& src_file, const std::string& type, const std::string& out_file) {
    if (type != "min" && type != "max" && type != "mean" && type != "direct") {
        std::cerr << "[gen-ascii-image] Type must be one of: min, max, mean, direct." << std::endl;
        print_usage();
        exit(1);
    }
    try {
        SATLAB lab(512, 512);
        cv::Mat directImg, minImg, maxImg, meanImg;
        lab.createByteImages(src_file, directImg, minImg, maxImg, meanImg);
        cv::Mat* outimg = nullptr;
        if (type == "min") outimg = &minImg;
        else if (type == "max") outimg = &maxImg;
        else if (type == "mean") outimg = &meanImg;
        else if (type == "direct") outimg = &directImg;
        cv::imwrite(out_file, *outimg);
        std::cout << "[gen-ascii-image] Processed: " << src_file << " -> " << out_file << std::endl;
    } catch (const std::exception& e) {
        std::cerr << "[gen-ascii-image] Error: " << e.what() << std::endl;
        exit(2);
    }
}

void run_gen_bin_scope_images(const std::string& src_file, const std::string& type, const std::string& out_dir) {
    if (type != "origin" && type != "sorted") {
        std::cerr << "[gen-bin-scope-images] Type must be 'origin' or 'sorted'." << std::endl;
        print_usage();
        exit(1);
    }
    try {
        SATLAB lab(512, 512);
        lab.parse(src_file.c_str());
        if (type == "sorted") {
            lab.mat.sortVecs();
        }
        lab.updateAllStats();
        lab.createAllBinScopeImages(out_dir);
        std::cout << "[gen-bin-scope-images] Processed: " << src_file << " -> " << out_dir << std::endl;
    } catch (const std::exception& e) {
        std::cerr << "[gen-bin-scope-images] Error: " << e.what() << std::endl;
        exit(2);
    }
}

void run_gen_block_images(const std::string& src_file, const std::string& type, const std::string& out_dir) {
    if (type != "origin" && type != "sorted") {
        std::cerr << "[gen-block-images] Type must be 'origin' or 'sorted'." << std::endl;
        print_usage();
        exit(1);
    }
    try {
        SATLAB lab(512, 512);
        lab.parse(src_file.c_str());
        if (type == "sorted") {
            lab.mat.sortVecs();
        }
        
        // Create output directory
        std::filesystem::create_directories(out_dir);
        
        // Process blocks row by row
        size_t totalBlockRows = (lab.claNum + lab.blockClauseSize - 1) / lab.blockClauseSize;
        
        for (size_t rowIndex = 0; rowIndex < totalBlockRows; rowIndex++) {
            // Initialize block row
            lab.initializeBlockRow(rowIndex);
            
            // Update stats for this block row
            lab.updateAllStatsForBlockRow(rowIndex);
            
            // Generate images for this block row
            lab.createCurrentBlockRowImages(out_dir, rowIndex);
            
            // Clear block row to free memory
            lab.clearBlockRow();
        }
        
        std::cout << "[gen-block-images] Processed: " << src_file << " -> " << out_dir << std::endl;
    } catch (const std::exception& e) {
        std::cerr << "[gen-block-images] Error: " << e.what() << std::endl;
        exit(2);
    }
}

int main(int argc, char* argv[]) {
    if (argc < 2) {
        print_usage();
        return 1;
    }
    std::string cmd = argv[1];
    if (cmd == "gen-scope-image") {
        if (argc != 5) {
            print_usage();
            return 1;
        }
        std::string src_file = argv[2];
        std::string type = argv[3];
        std::string out_file = argv[4];
        run_gen_scope_image(src_file, type, out_file);
    } else if (cmd == "gen-ascii-image") {
        if (argc != 5) {
            print_usage();
            return 1;
        }
        std::string src_file = argv[2];
        std::string type = argv[3];
        std::string out_file = argv[4];
        run_gen_ascii_image(src_file, type, out_file);
    } else if (cmd == "gen-bin-scope-images") {
        if (argc != 5) {
            print_usage();
            return 1;
        }
        std::string src_file = argv[2];
        std::string type = argv[3];
        std::string out_dir = argv[4];
        run_gen_bin_scope_images(src_file, type, out_dir);
    } else if (cmd == "gen-block-images") {
        if (argc != 5) {
            print_usage();
            return 1;
        }
        std::string src_file = argv[2];
        std::string type = argv[3];
        std::string out_dir = argv[4];
        run_gen_block_images(src_file, type, out_dir);
    } else {
        print_usage();
        return 1;
    }
    return 0;
}
